package cc.springai.openai.controller;

import com.openai.client.OpenAIClient;
import com.openai.client.okhttp.OpenAIOkHttpClient;
import com.openai.core.http.AsyncStreamResponse;
import com.openai.models.ChatModel;
import com.openai.models.chat.completions.ChatCompletion;
import com.openai.models.chat.completions.ChatCompletionChunk;
import com.openai.models.chat.completions.ChatCompletionCreateParams;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Optional;

@RestController
@RequestMapping("sdk")
public class SdkController {

    String apiKey;
    String baseUrl;

    public SdkController() {
        this.apiKey = System.getenv("API_KEY");
        this.baseUrl = System.getenv("BASE_URL");
    }

    @RequestMapping("chat")
    public String chat() {
        // 构建请求客户端
        OpenAIClient client = OpenAIOkHttpClient.builder()
                .baseUrl(baseUrl)
                .apiKey(apiKey)
                .build();

        // 构建请求参数
        ChatCompletionCreateParams params = ChatCompletionCreateParams.builder()
                .addUserMessage("Say this is a test")
                .model(ChatModel.GPT_4_1)
                .build();

        // 执行请求
        ChatCompletion chatCompletion = client.chat().completions().create(params);
        return chatCompletion.choices().stream()
                .findFirst()
                .map(choice -> choice.message().content())
                .get()
                .get();
    }

    @RequestMapping("stream")
    public SseEmitter stream() {
        // 构建请求客户端
        OpenAIClient client = OpenAIOkHttpClient.builder()
                .baseUrl(baseUrl)
                .apiKey(apiKey)
                .build();

        // 构建请求参数
        ChatCompletionCreateParams params = ChatCompletionCreateParams.builder()
                .addUserMessage("Tell a 100 word joke")
                .model(ChatModel.GPT_4_1)
                .build();

        SseEmitter emitter = new SseEmitter(Long.MAX_VALUE);

        client.async().chat().completions().createStreaming(params).subscribe(new AsyncStreamResponse.Handler<ChatCompletionChunk>() {
            @Override
            public void onNext(ChatCompletionChunk chunk) {
                System.out.println(chunk);
                try {
                    emitter.send(chunk);
                } catch (Exception e) {
                    e.printStackTrace();
                    emitter.completeWithError(e);
                }
            }

            @Override
            public void onComplete(Optional<Throwable> error) {
                if (error.isPresent()) {
                    try {
                        emitter.completeWithError(error.get()); // 发生错误时完成并抛出异常
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                } else {// 标记完成
                    emitter.complete(); // 完成SseEmitter
                }
            }
        });

        return emitter;
    }
}
